# code adapted from https://github.com/wendelinboehmer/dcg

import torch.nn as nn
import torch.nn.functional as F
import torch as th

class Discriminator(nn.Module):
    def __init__(self, state_dim, action_dim, args):
        super().__init__()
        self.dropout = args.dropout
        self.fc = nn.Sequential(nn.Linear(state_dim + action_dim, args.hidden_dim),
                               nn.ReLU(),
                               nn.Dropout(self.dropout),
                               nn.Linear(args.hidden_dim, args.hidden_dim),
                               nn.ReLU(),
                               nn.Dropout(self.dropout),
                               nn.Linear(args.hidden_dim, args.hidden_dim),
                               nn.ReLU(),                               
                               nn.Dropout(self.dropout),
                               nn.Linear(args.hidden_dim, 1)
                               )
    
    def forward(self, state, action):
        x = th.cat([state, action], dim=-1)
        return self.fc(x).sigmoid()

    
class GAILAgent(nn.Module):
    def __init__(self, input_shape, args):
        super(GAILAgent, self).__init__()
        self.args = args

        self.fc1 = nn.Linear(input_shape, args.hidden_dim)
        if self.args.use_rnn:
            self.rnn = nn.GRUCell(args.hidden_dim, args.hidden_dim)
        else:
            self.rnn = nn.Linear(args.hidden_dim, args.hidden_dim)
        self.fc2 = nn.Linear(args.hidden_dim, args.n_actions)
        self.D = nn.ModuleList([Discriminator(args.state_shape, args.n_actions, args) for _ in range(args.n_agents)])   

    def init_hidden(self):
        # make hidden states on same device as model
        return self.fc1.weight.new(1, self.args.hidden_dim).zero_()

    def forward(self, inputs, hidden_state):
        x = F.relu(self.fc1(inputs))
        h_in = hidden_state.reshape(-1, self.args.hidden_dim)
        if self.args.use_rnn:
            h = self.rnn(x, h_in)
        else:
            h = F.relu(self.rnn(x))
        q = self.fc2(h)
        return q, h

